import logging
from functools import partial

import numpy as np
import torch

from datasets.patch_pair_dataset import PatchPairDataset
from datasets.patch_clipper import PatchClipper

logger = logging.getLogger("match")

class TemplateMatchDataset(PatchPairDataset, PatchClipper):


    def __init__(self, config=None, *, tools=None, name=""):
        super(TemplateMatchDataset, self).__init__(config, super=True)
        super(PatchPairDataset, self).__init__(**(config.clipper.toDict()))
            
        #-- dataset
        self.dt = tools(config.base_dir, config.list_file)
        self.pl = self.dt.get_pair_list()

        logger.info(f"{name} Patch Pair Dataset created with {self.__len__()} pairs ---- search domain is {self.search_domain}")


    def __getitem__(self, index):
        """
        ((sar, opt), heat map)
        """
        (sar_tensor, opt_tensor), _ = super().__getitem__(index)

        sar_t, opt_t, gt_map, shift = self.crop(sar_tensor, opt_tensor)

        return (sar_t, opt_t), (gt_map, shift)


class MaskTemplateMatchDataset(TemplateMatchDataset):


    def __init__(self, config=None, *, tools=None, name=""):
        super(MaskTemplateMatchDataset, self).__init__(
            config, tools=tools, name=name
        )


    def __getitem__(self, index):
        #-- augmentor
        if self.augmentor:
            self.augmentor.refresh_random_state()

        img_sar, img_opt, labels, mask = self._load_data(index)

        if self.augmentor is not None:
            img_sar = self.augmentor(img_sar)
            img_opt = self.augmentor(img_opt)
            mask = self.augmentor(mask)

        
        img_sar = img_sar[:, :, 0]
        img_opt = img_opt[:, :, 0]
        mask = mask[:, :, 0]
        img_sar = np.ascontiguousarray(img_sar)
        img_opt = np.ascontiguousarray(img_opt)
        mask = np.ascontiguousarray(mask)
        img_sar = (img_sar - img_sar.min())/(img_sar.ptp())
        img_opt = (img_opt - img_opt.min())/(img_opt.ptp())
        
        img_opt = np.expand_dims(img_opt, axis=2)
        img_sar = np.expand_dims(img_sar, axis=2)
        mask = np.expand_dims(mask, axis=2)
        
        img_sar_tensor = self.transforms(img_sar).float()
        img_opt_tensor = self.transforms(img_opt).float()
        mask_tensor = self.transforms(mask).float()

        labels = torch.tensor(labels)

        sar_t, opt_t, gt_map, shift = self.crop(img_sar_tensor, img_opt_tensor, mask=mask_tensor)

        return (sar_t[0], opt_t, sar_t[1]), (gt_map, shift)


    def _load_data(self, index, drop=False):
        '''
        index to transformed image pairs and other info 
        form caches if exists or files

        Return:
        ------
            {"SAR" : sar patch, 
            "OPT" : opt patch, 
            "LABEL" : labels, 
            "MASK" : mask
            }
        '''

        data = self._try_cache(index)

        cache_key = " "
        
        if not data:
            data = {"SAR": None, "OPT": None, "Y": None, "INFO": None}

        while (data["OPT"] is None) and (data["SAR"] is None):
            pair = {
                'OPT': self.pl[index][0], 
                'SAR': self.pl[index][1], 
                'MASK': self.pl[index][2]
            }

            imgs = self._get_patches(pair, 
                opt_transform=self._normalize_scale, 
                sar_transform=self._normalize_scale, 
            )

            if len(imgs["OPT"]) > 0 and len(imgs["SAR"]) > 0:
                data["OPT"] = imgs["OPT"][0]
                data["SAR"] = imgs["SAR"][0]
                data["INFO"] = imgs["MASK"][0]
                data["LABEL"] = 0
                
                if self.cache is not None:
                    self.cache[cache_key] = data

            else:
                if drop:
                    break

                # HACK: you shouldn't come here
                index += 1

        return  data["SAR"], data["OPT"], data["LABEL"], data["INFO"]


    def _get_patches(
        self, pair: dict, opt_transform=None, sar_transform=None, **kwargs
    ):
        """
        open the image patch, transform and return
        """
        imgs = {"OPT": [], "SAR": [], "MASK": []}

        for modal, filename in pair.items():
            transform = sar_transform if modal == "SAR" else opt_transform
            im = self.dt.get_patch(
                filename, transform=transform, modal=modal
            )

            if im is not None:
                im = im[:,:]
                imgs[modal].append(im)
                
        return imgs


if __name__ == "__main__":
    import cv2 as cv
    from dotmap import DotMap

    from tools.dataset_tools import pm_mask_tools

    class config:
        def __init__(self):
            self.single_domain = ""
            self.cache_dir = None
            self.cache_size = 0
            self.base_dir = None
            self.list_file = "E:/workspace/SOMatch/json/os-select/1-s-test-mask.json"
            self.augment = True
            self.clipper = DotMap({
                "name": "shift template",
                "template_size": 144,
                "shift_x": 6,
                "shift_y": 6,
                "search_rad": 32,
                "search_domain": "B", 
                "fixed": True
            })


    cfg = config()

    ppd = MaskTemplateMatchDataset(cfg, tools=pm_mask_tools, name="OS")

    (img0, img1, mask), _ = ppd[0]
    # img0, img1, gt_map = ppd._crop_patchs(img0, img1)

    img0 = img0.numpy()
    img1 = img1.numpy()
    mask = mask.numpy()
    img0 = np.transpose(img0, [1, 2, 0])
    img1 = np.transpose(img1, [1, 2, 0])
    mask = np.transpose(mask, [1, 2, 0])

    cv.imwrite("./tmp/image/dataset_test/img0.png", img0*255)
    cv.imwrite("./tmp/image/dataset_test/img1.png", img1*255)
    cv.imwrite("./tmp/image/dataset_test/mask.png", mask*255)

